package de.unibi.comet.ac;

import java.util.List;
import java.util.ArrayList;
import de.unibi.comet.hmm.Emission;

/** Distribution of integer masses. Internally uses an ArrayList of
 *  size "maximum mass+1". Hence it is space consuming for large
 *  masses. */
public class MassDistribution {
	private ArrayList<Double> dist;

	/** Creates an empty distribution. */
	public MassDistribution() {
		dist = new ArrayList<Double>();
	}

	/** Create a "delta"-distribution with p(mass)=1. */
	public MassDistribution(int mass) {
		dist = new ArrayList<Double>(mass+1);
		growTo(mass-1);
		dist.add(1.0);
	}

	private void growTo(int mass) {
		dist.ensureCapacity(mass+2);
		for (int i=dist.size(); i<=mass; ++i) {
			dist.add(0.0);
		}
	}
	
	public double get(int mass) {
		if ((mass<0) || (mass>=dist.size())) {
			return 0.0;
		}
		return dist.get(mass);
	}
	
	public void set(int mass, double prob) {
		if (mass<0) throw new IllegalArgumentException("Masses <0 are not supported");
		if (prob<0.0) throw new IllegalArgumentException("Probabilities must be >= 0.0");
		if (mass>=dist.size()) {
			growTo(mass-1);
			dist.add(prob);
		} else {
			dist.set(mass, prob);
		}
	}
	
	public List<Emission> toEmissionList() {
		return toEmissionList(0);
	}
	
	public List<Emission> toEmissionList(int massOffset) {
		List<Emission> result = new ArrayList<Emission>();
		for (int i=0; i<dist.size(); ++i) {
			if (dist.get(i)>0.0) result.add(new Emission(i+massOffset, dist.get(i)));
		}
		return result;
	}
	
	/** Multiplies every entry by given factor. 
	 *  NOTE: Normalization (sum=1.0) is not ensured automatically,
	 *    use method normalize for this.  */
	public void multiply(double factor) {
		for (int i=0; i<dist.size(); ++i) {
			dist.set(i, dist.get(i)*factor);
		}
	}

	/** Adds a specific value to given mass. */
	public void add(int mass, double value) {
		if (mass>=dist.size()) {
			growTo(mass-1);
			dist.add(value);
		} else {
			dist.set(mass, dist.get(mass)+value);
		}
	}
	
	/** Adds another distribution to this one. 
	 *  NOTE: Normalization (sum=1.0) is not ensured automatically,
	 *    use method normalize for this.  */
	public void add(MassDistribution summand) {
		growTo(summand.dist.size()-1);
		for (int i=0; i<summand.dist.size(); ++i) {
			dist.set(i, dist.get(i)+summand.dist.get(i));
		}
	}
	
	/** Forces sum over all probabilities to be one.*/
	public void normalize() {
		double sum = 0.0;
		for (double d : dist) sum+=d;
		if (sum<1e-12) throw new IllegalStateException("Cannot normalize (sum is zero)");
		multiply(sum);
	}
	
	/** Converts current distribution into the distribution of accumulated masses. 
	 *  (The distribution that would be observed when drawing a mass from each of 
	 *  the two distributions and add them up.) */
	public void join(MassDistribution distribution) {
		int n = dist.size()+distribution.dist.size()-1;
		ArrayList<Double> newDist = new ArrayList<Double>(n);
		for (int i=0; i<n; ++i) newDist.add(0.0);
		
		for (int i=0; i<dist.size(); ++i) {
			for (int j=0; j<distribution.dist.size(); ++j) {
				newDist.set(i+j, newDist.get(i+j)+dist.get(i)*distribution.dist.get(j));
			}
		}
		dist = newDist;
	}
}
